/* Copyright (C) 2010 Luca Piccinelli
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/*
 * File:   normal_bayes_model.cpp
 * Author: picci
 *
 * Created on December 4, 2010, 11:57 AM
 */

#include <algorithm>

#include "machine_learning/statistical_models/normal_bayes_model.h"
#include "generators/genericGenerator.h"
#include "factory/factories.h"

using namespace cv;
using namespace std;

namespace NAMESPACE{

// ****** Abstract class implementation ****************************************
AbstractNormalBayesModel::~AbstractNormalBayesModel(){}
AbstractNormalBayesModel::AbstractNormalBayesModel() : NormalBayesModel(){}
AbstractNormalBayesModel::AbstractNormalBayesModel(const AbstractNormalBayesModel& nbm) : NormalBayesModel(){}
AbstractNormalBayesModel& AbstractNormalBayesModel::operator= (const AbstractNormalBayesModel& nbm){
    if(this == &nbm) return *this;
    return *this;
}
// -----------------------------------------------------------------------------

// ****** Discret2D model definition *********************************************:
//  _____________________________________________________
// |   Standard constructors                             |
// |_____________________________________________________|

Discret2DNormalBayesModel::~Discret2DNormalBayesModel(){
    delete a_priori_data_lut;
    delete a_priori_evidence;
    delete with_evidence_data_lut;

    delete with_evidence_data_lut_counters;
    delete a_priori_data_lut_counters;
}
Discret2DNormalBayesModel::Discret2DNormalBayesModel() : AbstractNormalBayesModel(),
                                                         with_evidence_data_lut(new vector<Mat>()),
                                                         a_priori_data_lut(new Mat()),
                                                         a_priori_evidence(new Mat()),
                                                         range_begin(),
                                                         range_end(),
                                                         tot_train_counter(0),
                                                         class_train_counters(),
                                                         with_evidence_data_lut_counters(new vector<Mat>()),
                                                         a_priori_data_lut_counters(new Mat()){}
Discret2DNormalBayesModel::Discret2DNormalBayesModel(const Discret2DNormalBayesModel& nbm)
                                                    : AbstractNormalBayesModel(nbm),
                                                      with_evidence_data_lut(new vector<Mat>(*(nbm.with_evidence_data_lut))),
                                                      a_priori_data_lut(new Mat(*(nbm.a_priori_data_lut))),
                                                      a_priori_evidence(new Mat(*(nbm.a_priori_evidence))),
                                                      range_begin(nbm.range_begin),
                                                      range_end(nbm.range_end),
                                                      tot_train_counter(nbm.tot_train_counter),
                                                      class_train_counters(nbm.class_train_counters),
                                                      with_evidence_data_lut_counters(new vector<Mat>(*(nbm.with_evidence_data_lut))),
                                                      a_priori_data_lut_counters(new Mat(*(nbm.a_priori_data_lut))){}
Discret2DNormalBayesModel& Discret2DNormalBayesModel::operator= (const Discret2DNormalBayesModel& nbm){
    if(this == &nbm) return *this;
    AbstractNormalBayesModel::operator =(nbm);

    *with_evidence_data_lut = *(nbm.with_evidence_data_lut);
    *a_priori_data_lut      = *(nbm.a_priori_data_lut);
    *a_priori_evidence      = *(nbm.a_priori_evidence);
    range_begin             = nbm.range_begin;
    range_end               = nbm.range_end;
    tot_train_counter       = nbm.tot_train_counter;
    class_train_counters    = nbm.class_train_counters;

    *with_evidence_data_lut_counters = *(nbm.with_evidence_data_lut_counters);
    *a_priori_data_lut_counters      = *(nbm.a_priori_data_lut_counters);

    return *this;
}
// .......................................................

//  _____________________________________________________
// |   Other constructors                                |
// |_____________________________________________________|

Discret2DNormalBayesModel::Discret2DNormalBayesModel(const Point2i& _range_begin, const Point2i& _range_end)
                                                    : AbstractNormalBayesModel(),
                                                      with_evidence_data_lut(new vector<Mat>()),
                                                      a_priori_data_lut(new Mat()),
                                                      a_priori_evidence(new Mat()),
                                                      range_begin(_range_begin),
                                                      range_end(_range_end),
                                                      tot_train_counter(0),
                                                      class_train_counters(),
                                                      with_evidence_data_lut_counters(new vector<Mat>()),
                                                      a_priori_data_lut_counters(new Mat()){}

Discret2DNormalBayesModel::Discret2DNormalBayesModel(int bx, int by, int ex, int ey)
                                                    : AbstractNormalBayesModel(),
                                                      with_evidence_data_lut(new vector<Mat>()),
                                                      a_priori_data_lut(new Mat()),
                                                      a_priori_evidence(new Mat()),
                                                      range_begin(bx, by),
                                                      range_end(ex, ey),
                                                      tot_train_counter(0),
                                                      class_train_counters(),
                                                      with_evidence_data_lut_counters(new vector<Mat>()),
                                                      a_priori_data_lut_counters(new Mat()){}
// .......................................................

NormalBayesModel& Discret2DNormalBayesModel::train(const Mat& train_data, const Mat& responses, bool update){
    int dimension = 2;

    CV_Assert( train_data.cols    == dimension           );
    CV_Assert( responses.cols     == 1                   );
    CV_Assert( train_data.rows    == responses.rows      );

    uint classes_num = *(std::max_element( responses.begin<int>(), responses.end<int>() )) + 1;
    int cols_num = range_end.x - range_begin.x;
    int rows_num = range_end.y - range_begin.y;

    // Allocating memory if nedeed
    uint _class;
    if(!update || !with_evidence_data_lut_counters->size() || with_evidence_data_lut_counters->size() != classes_num){
        with_evidence_data_lut_counters->clear();
        for(_class = 0; _class < classes_num; _class++){
            with_evidence_data_lut_counters->push_back(Mat::zeros(rows_num, cols_num, CV_64F));
        }

        with_evidence_data_lut->resize(classes_num);

        a_priori_evidence->create(1, classes_num, CV_64F);
        class_train_counters.resize(classes_num, 0);

        *a_priori_data_lut_counters = Mat::zeros(rows_num, cols_num, CV_64F);

        tot_train_counter = 0;
    }

    tot_train_counter += train_data.rows;

    // Computing a priori probabilities and counting a posteriori occurrences
    int xy[2];
    double max_x_range_value = range_end.x - range_begin.x;
    double max_y_range_value = range_end.y - range_begin.y;
    for(int i = 0; i < train_data.rows; i++){
        xy[0] = train_data.at<int>(i, 0) - range_begin.x;
        xy[1] = train_data.at<int>(i, 1) - range_begin.y;
        if(xy[0] > max_x_range_value) xy[0] = max_x_range_value;
        else if(xy[0] < 0)            xy[0] = 0;
        if(xy[1] > max_y_range_value) xy[1] = max_y_range_value;
        else if(xy[1] < 0)            xy[1] = 0;
        _class = responses.at<int>(i, 0);

        with_evidence_data_lut_counters->at(_class).at<prob_t>(xy) += 1.0;
        a_priori_data_lut_counters->at<prob_t>(xy)                 += 1.0;
        class_train_counters[_class] += 1;
    }

    // Normalizing all counters
    vector<Mat>::iterator it;
    prob_t tot_scale_value = 1.0 / tot_train_counter;
    prob_t scale_value;

    it = with_evidence_data_lut_counters->begin();
    for(_class = 0; it != with_evidence_data_lut_counters->end(); it++, _class++) {
        scale_value = 1.0 / class_train_counters[_class];
        it->convertTo((*with_evidence_data_lut)[_class], CV_64F, scale_value);

        a_priori_evidence->at<prob_t>(_class) = class_train_counters[_class] * tot_scale_value;
    }
    a_priori_data_lut_counters->convertTo(*a_priori_data_lut, CV_64F, tot_scale_value);
    return *this;
}

void Discret2DNormalBayesModel::predict(const Mat& samples, Mat& responses) const{
    CV_Assert( samples.type() == DataType<int>::type &&
               samples.cols   == 2                     );

    int classes_num = a_priori_evidence->cols;
    Mat tmp_responses(classes_num, samples.rows, DataType<prob_t>::type);

    prob_t a_priori_evidence_prob, a_priori_data_scale, with_evidence_data_prob;
    for(int i = 0; i < samples.rows; i++){
        const int* r = samples.ptr<int>(i);
        a_priori_data_scale = 1.0 / a_priori_data_lut->at<prob_t>(r);
        for(int j = 0; j < classes_num; j++){
            a_priori_evidence_prob  = a_priori_evidence->at<prob_t>(0, j);
            with_evidence_data_prob = (*with_evidence_data_lut)[j].at<prob_t>(r);

            // bayes classifier rule
            tmp_responses.at<prob_t>(j, i) =  a_priori_evidence_prob * with_evidence_data_prob * a_priori_data_scale;
        }
    }    
    responses = tmp_responses;
}

std::vector<cv::Mat>& Discret2DNormalBayesModel::get_with_evidence_data_luts() const {
    return *with_evidence_data_lut;
}

cv::Mat& Discret2DNormalBayesModel::get_a_priori_data_lut() const {
    return *a_priori_data_lut;
}
cv::Mat& Discret2DNormalBayesModel::get_a_priori_evidence() const {
    return *a_priori_evidence;
}

Discret2DNormalBayesModel& Discret2DNormalBayesModel::set_with_evidence_data_luts(std::vector<cv::Mat>* luts){
    if(with_evidence_data_lut) delete with_evidence_data_lut;
    with_evidence_data_lut = luts;
    return *this;
}

bool Discret2DNormalBayesModel::is_trained() const{
    return with_evidence_data_lut->size();
}

NormalBayesModel& Discret2DNormalBayesModel::swap(NormalBayesModel& m){
    Discret2DNormalBayesModel* d2dm = dynamic_cast<Discret2DNormalBayesModel*>(&m);
    with_evidence_data_lut->swap(*(d2dm->with_evidence_data_lut));
    std::swap(a_priori_data_lut, d2dm->a_priori_data_lut);
    std::swap(a_priori_evidence, d2dm->a_priori_evidence);
    return *this;
}
// -----------------------------------------------------------------------------

// ****** Continuous model definition ******************************************
ContinuousNormalBayesModel::~ContinuousNormalBayesModel(){}
ContinuousNormalBayesModel::ContinuousNormalBayesModel() : AbstractNormalBayesModel(){}
ContinuousNormalBayesModel::ContinuousNormalBayesModel(const ContinuousNormalBayesModel& nbm) : AbstractNormalBayesModel(nbm){}
ContinuousNormalBayesModel& ContinuousNormalBayesModel::operator= (const ContinuousNormalBayesModel& nbm){
    if(this == &nbm) return *this;
    AbstractNormalBayesModel::operator =(nbm);
    return *this;
}

NormalBayesModel& ContinuousNormalBayesModel::train(const Mat& train_data, const Mat& responses, bool update){ throw std::runtime_error("Not implemented"); }
void ContinuousNormalBayesModel::predict(const Mat& samples, Mat& responses) const{ throw std::runtime_error("Not implemented"); }

bool ContinuousNormalBayesModel::is_trained() const{ throw std::runtime_error("Not implemented"); }
NormalBayesModel& ContinuousNormalBayesModel::swap(NormalBayesModel& m){ throw std::runtime_error("Not implemented"); }
// -----------------------------------------------------------------------------


typedef GenericGenerator<Discret2DNormalBayesModel, NormalBayesModel> GenericGeneratorDiscret2DNormalBayesModel;
REGISTER_PARSE_TYPE( GenericGeneratorDiscret2DNormalBayesModel );
GenericGeneratorFactory<Discret2DNormalBayesModel, NormalBayesModel> instance_of_D2DNBMGenericGeneratorFactory;

}
